package connpool

import (
	"context"
	"database/sql"
	"gorm.io/gorm"
	"sync"
)

// WriteSplit 读写分离
type WriteSplit struct {
	master gorm.ConnPool
	slaves []SlavePool
	mutex  sync.Mutex
}

func (w WriteSplit) PrepareContext(ctx context.Context, query string) (*sql.Stmt, error) {
	return w.master.PrepareContext(ctx, query)
}

func (w WriteSplit) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
	return w.master.ExecContext(ctx, query, args...)
}

func (w WriteSplit) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
	w.mutex.Lock()
	w.mutex.Unlock()

	//在从库上读
	var total int
	var maxSlve SlavePool
	for _, slave := range w.slaves {
		if !slave.available {
			continue
		}

		total += slave.weight
		slave.currentWeight += slave.weight
		if maxSlve.ConnPool == nil || slave.currentWeight > maxSlve.currentWeight {
			maxSlve = slave
		}
	}
	maxSlve.currentWeight -= total
	return maxSlve.QueryContext(ctx, query, args...)
}

func (w WriteSplit) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
	w.mutex.Lock()
	w.mutex.Unlock()

	//在从库上读
	var total int
	var maxSlve SlavePool
	for _, slave := range w.slaves {
		if !slave.available {
			continue
		}

		total += slave.weight
		slave.currentWeight += slave.weight
		if maxSlve.ConnPool == nil || slave.currentWeight > maxSlve.currentWeight {
			maxSlve = slave
		}
	}
	maxSlve.currentWeight -= total
	return maxSlve.QueryRowContext(ctx, query, args...)
}

type SlavePool struct {
	gorm.ConnPool
	//初始权重
	weight int
	//当前权重
	currentWeight int
	available     bool
}
